{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VQModel(\n",
      "  (encoder): Encoder(\n",
      "    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (down_blocks): ModuleList(\n",
      "      (0): DownEncoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "        (downsamplers): ModuleList(\n",
      "          (0): Downsample2D(\n",
      "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (1): DownEncoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "            (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "        (downsamplers): ModuleList(\n",
      "          (0): Downsample2D(\n",
      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (2): DownEncoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "            (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (mid_block): UNetMidBlock2D(\n",
      "      (attentions): ModuleList(\n",
      "        (0): AttentionBlock(\n",
      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (query): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (key): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (value): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "      )\n",
      "      (resnets): ModuleList(\n",
      "        (0): ResnetBlock2D(\n",
      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (dropout): Dropout(p=0.0, inplace=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (nonlinearity): SiLU()\n",
      "        )\n",
      "        (1): ResnetBlock2D(\n",
      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (dropout): Dropout(p=0.0, inplace=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (nonlinearity): SiLU()\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "    (conv_act): SiLU()\n",
      "    (conv_out): Conv2d(512, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  )\n",
      "  (quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n",
      "  (quantize): VectorQuantizer(\n",
      "    (embedding): Embedding(8192, 3)\n",
      "  )\n",
      "  (post_quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n",
      "  (decoder): Decoder(\n",
      "    (conv_in): Conv2d(3, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (up_blocks): ModuleList(\n",
      "      (0): UpDecoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "          (2): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "        (upsamplers): ModuleList(\n",
      "          (0): Upsample2D(\n",
      "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (1): UpDecoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "            (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "          (2): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "        (upsamplers): ModuleList(\n",
      "          (0): Upsample2D(\n",
      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (2): UpDecoderBlock2D(\n",
      "        (resnets): ModuleList(\n",
      "          (0): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "            (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "          )\n",
      "          (1): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "          (2): ResnetBlock2D(\n",
      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "            (nonlinearity): SiLU()\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (mid_block): UNetMidBlock2D(\n",
      "      (attentions): ModuleList(\n",
      "        (0): AttentionBlock(\n",
      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (query): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (key): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (value): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "      )\n",
      "      (resnets): ModuleList(\n",
      "        (0): ResnetBlock2D(\n",
      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (dropout): Dropout(p=0.0, inplace=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (nonlinearity): SiLU()\n",
      "        )\n",
      "        (1): ResnetBlock2D(\n",
      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
      "          (dropout): Dropout(p=0.0, inplace=False)\n",
      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (nonlinearity): SiLU()\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
      "    (conv_act): SiLU()\n",
      "    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from yldiffusers import unconditional_garbage_vqvae_config_dict\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from yldiffusers import unconditional_vqvae_config_dict, get_config, VAEConfig\n",
    "config = get_config(unconditional_garbage_vqvae_config_dict, VAEConfig)\n",
    "print(config.vae)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "🚀yl-diffusion-VQVAE training starts!\n"
     ]
    },
    {
     "data": {
      "text/plain": "HBox(children=(HTML(value='train Epoch [1/4]'), FloatProgress(value=0.0, max=5923.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c2f9b527f0a64b41ab81ddf7b556d0c6"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "\u001B[1;32m<ipython-input-2-398df57bd005>\u001B[0m in \u001B[0;36m<module>\u001B[1;34m\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[1;32mfrom\u001B[0m \u001B[0myldiffusers\u001B[0m \u001B[1;32mimport\u001B[0m \u001B[0mVQVAE\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m      2\u001B[0m \u001B[0mvqvae\u001B[0m \u001B[1;33m=\u001B[0m \u001B[0mVQVAE\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mconfig\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m----> 3\u001B[1;33m \u001B[0mvqvae\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mtrain\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m",
      "\u001B[1;32mD:\\毕业论文\\ylcv\\diffusion-pytorch\\yldiffusers\\models\\models.py\u001B[0m in \u001B[0;36mtrain\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m    450\u001B[0m                     \u001B[0mtotal_loss\u001B[0m \u001B[1;33m=\u001B[0m  \u001B[0mvqvae_loss\u001B[0m \u001B[1;33m+\u001B[0m \u001B[0mdecode_loss\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m    451\u001B[0m                     \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0maccelerator\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mbackward\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mtotal_loss\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m--> 452\u001B[1;33m                     \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mlogger\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;34m'train_vqvae_loss'\u001B[0m\u001B[1;33m]\u001B[0m \u001B[1;33m+=\u001B[0m \u001B[0mvqvae_loss\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mdetach\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mitem\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m    453\u001B[0m                     \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mlogger\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;34m'train_decode_loss'\u001B[0m\u001B[1;33m]\u001B[0m \u001B[1;33m+=\u001B[0m \u001B[0mdecode_loss\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mdetach\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mitem\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m    454\u001B[0m                     \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mlogger\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;34m'train_total_num'\u001B[0m\u001B[1;33m]\u001B[0m \u001B[1;33m+=\u001B[0m \u001B[0mclean_images\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;36m0\u001B[0m\u001B[1;33m]\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "from yldiffusers import VQVAE\n",
    "vqvae = VQVAE(config)\n",
    "vqvae.train()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "from yldiffusers import VQVAE\n",
    "vqvae = VQVAE()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "vqvae.load('unconditional_vqvae/exp_0/VQModel.pt', 'unconditional_vqvae/exp_0/unconditional_vqvae_config.pickle')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5 epoch eval has done!\n"
     ]
    }
   ],
   "source": [
    "vqvae.eval()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "data": {
      "text/plain": "VQModel(\n  (encoder): Encoder(\n    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (down_blocks): ModuleList(\n      (0): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (downsamplers): ModuleList(\n          (0): Downsample2D(\n            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n          )\n        )\n      )\n      (1): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (downsamplers): ModuleList(\n          (0): Downsample2D(\n            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n          )\n        )\n      )\n      (2): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n      )\n    )\n    (mid_block): UNetMidBlock2D(\n      (attentions): ModuleList(\n        (0): AttentionBlock(\n          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (query): Linear(in_features=512, out_features=512, bias=True)\n          (key): Linear(in_features=512, out_features=512, bias=True)\n          (value): Linear(in_features=512, out_features=512, bias=True)\n          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n        )\n      )\n      (resnets): ModuleList(\n        (0): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n        (1): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n      )\n    )\n    (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n    (conv_act): SiLU()\n    (conv_out): Conv2d(512, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n  )\n  (quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n  (quantize): VectorQuantizer(\n    (embedding): Embedding(8192, 3)\n  )\n  (post_quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n  (decoder): Decoder(\n    (conv_in): Conv2d(3, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (up_blocks): ModuleList(\n      (0): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (upsamplers): ModuleList(\n          (0): Upsample2D(\n            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n      )\n      (1): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (upsamplers): ModuleList(\n          (0): Upsample2D(\n            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n      )\n      (2): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n      )\n    )\n    (mid_block): UNetMidBlock2D(\n      (attentions): ModuleList(\n        (0): AttentionBlock(\n          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (query): Linear(in_features=512, out_features=512, bias=True)\n          (key): Linear(in_features=512, out_features=512, bias=True)\n          (value): Linear(in_features=512, out_features=512, bias=True)\n          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n        )\n      )\n      (resnets): ModuleList(\n        (0): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n        (1): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n      )\n    )\n    (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n    (conv_act): SiLU()\n    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n  )\n)"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vqvae.vae"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "data": {
      "text/plain": "VQModel(\n  (encoder): Encoder(\n    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (down_blocks): ModuleList(\n      (0): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (downsamplers): ModuleList(\n          (0): Downsample2D(\n            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n          )\n        )\n      )\n      (1): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (downsamplers): ModuleList(\n          (0): Downsample2D(\n            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n          )\n        )\n      )\n      (2): DownEncoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n      )\n    )\n    (mid_block): UNetMidBlock2D(\n      (attentions): ModuleList(\n        (0): AttentionBlock(\n          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (query): Linear(in_features=512, out_features=512, bias=True)\n          (key): Linear(in_features=512, out_features=512, bias=True)\n          (value): Linear(in_features=512, out_features=512, bias=True)\n          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n        )\n      )\n      (resnets): ModuleList(\n        (0): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n        (1): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n      )\n    )\n    (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n    (conv_act): SiLU()\n    (conv_out): Conv2d(512, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n  )\n  (quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n  (quantize): VectorQuantizer(\n    (embedding): Embedding(8192, 3)\n  )\n  (post_quant_conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n  (decoder): Decoder(\n    (conv_in): Conv2d(3, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (up_blocks): ModuleList(\n      (0): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (upsamplers): ModuleList(\n          (0): Upsample2D(\n            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n      )\n      (1): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n            (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n        (upsamplers): ModuleList(\n          (0): Upsample2D(\n            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n      )\n      (2): UpDecoderBlock2D(\n        (resnets): ModuleList(\n          (0): ResnetBlock2D(\n            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n            (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n            (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n          )\n          (1): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n          (2): ResnetBlock2D(\n            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n            (dropout): Dropout(p=0.0, inplace=False)\n            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            (nonlinearity): SiLU()\n          )\n        )\n      )\n    )\n    (mid_block): UNetMidBlock2D(\n      (attentions): ModuleList(\n        (0): AttentionBlock(\n          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (query): Linear(in_features=512, out_features=512, bias=True)\n          (key): Linear(in_features=512, out_features=512, bias=True)\n          (value): Linear(in_features=512, out_features=512, bias=True)\n          (proj_attn): Linear(in_features=512, out_features=512, bias=True)\n        )\n      )\n      (resnets): ModuleList(\n        (0): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n        (1): ResnetBlock2D(\n          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n          (dropout): Dropout(p=0.0, inplace=False)\n          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (nonlinearity): SiLU()\n        )\n      )\n    )\n    (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n    (conv_act): SiLU()\n    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n  )\n)"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vqvae.vae"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "index = np.random.randint(0,100,3)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "67\n",
      "50\n",
      "47\n"
     ]
    }
   ],
   "source": [
    "for i in index:\n",
    "    print(i)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}